import os
import sys
import xml.dom.minidom as xmldom

import cv2
import numpy as np

from tool.MyTool import remove_salt

sys.path.append('..')


class template:
    '''
    模板类
    '''
    tempLabel = ""
    x = -1
    y = -1
    width = 0
    height = 0
    img = None

    def __init__(self, label, x, y, width, height, img):
        self.tempLabel = label
        self.x = x
        self.y = y
        self.width = width
        self.height = height
        self.img = img

    def area(self):
        return self.width * self.height


class match_result:
    '''
    匹配结果
    '''

    pt = None
    template = None
    width = None
    height = None
    area = None

    def __init__(self, pt, template):
        self.pt = pt
        self.template = template
        self.area = template.area()
        self.width = template.width
        self.height = template.height


def filter_out_black(src_frame):
    '''
    提取黑色
    :param src_frame:
    :return:
    '''
    if src_frame is not None:
        hsv = cv2.cvtColor(src_frame, cv2.COLOR_BGR2HSV)
        lower_val = np.array([0, 0, 0])
        upper_val = np.array([179, 255, 127])
        mask = cv2.inRange(hsv, lower_val, upper_val)
        mask_inv = cv2.bitwise_not(mask)
        return mask_inv


def read_templates(rootPath):
    '''
    读取所有的模板
    :param rootPath:
    :return:
    '''
    files = os.listdir(rootPath)

    all_template = []
    template_map = {}

    for file in files:
        if file.endswith(".xml"):
            xmlfname = os.path.join(rootPath, file)
            fname = os.path.join(rootPath, file.replace(".xml", ".png"))
            domobj = xmldom.parse(xmlfname)
            elementobjs = domobj.documentElement
            subobjs = elementobjs.getElementsByTagName("object")

            img = cv2.imread(fname)
            cv2.threshold(img, 180, 255, cv2.THRESH_BINARY, img)
            img = filter_out_black(img)

            img = remove_salt(img, 0, 255, 255, 8)

            cv2.imshow("ig", img)

            cv2.waitKey(10000)

            for subobj in subobjs:
                name = subobj.getElementsByTagName("name")[0].firstChild.data
                bndbox = subobj.getElementsByTagName("bndbox")[0]
                xmin = int(bndbox.getElementsByTagName("xmin")[0].firstChild.data)
                ymin = int(bndbox.getElementsByTagName("ymin")[0].firstChild.data)
                xmax = int(bndbox.getElementsByTagName("xmax")[0].firstChild.data)
                ymax = int(bndbox.getElementsByTagName("ymax")[0].firstChild.data)

                template_img = img[ymin:ymax, xmin:xmax]  # 裁剪坐标为[y0:y1, x0:x1]

                tmp = template(label=name, x=xmin, y=ymin, width=(xmax - xmin), height=(ymax - ymin),
                               img=template_img)
                all_template.append(tmp)

                if template_map.__contains__(tmp.tempLabel):
                    template_map[tmp.tempLabel].append(tmp)
                else:
                    temp_list = []
                    temp_list.append(tmp)
                    template_map[tmp.tempLabel] = temp_list

    return template_map


def do_match(template_map, input_img, threshold):
    '''
    :param template_map: 所有的模板
    :param input_img:输入图像
    :return:
    '''

    match_results = []

    for type_label in template_map:
        type_label_list = template_map[type_label]
        for tmp_late in type_label_list:  # 对应模板所有数据都来匹配一波
            res = cv2.matchTemplate(input_img, tmp_late.img, cv2.TM_CCOEFF_NORMED)
            loc = np.where(res >= threshold)
            template_width = tmp_late.width
            template_height = tmp_late.height

            match_count = 0

            for pt in zip(*loc[::-1]):  # 展开结果

                bottom_right = (pt[0] + template_width, pt[1] + template_height)
                top_left = pt
                match_count = match_count + 1
                match_ret = match_result(pt, tmp_late)
                match_results.append(match_ret)

            # if 0 < match_count:
            #     print(tmp_late.tempLabel + "匹配到个数" + str(match_count))
            # print(res)

    def sort_key(tem):
        return tem.pt[0]

    match_results.sort(key=sort_key)
    return match_results


def classfy_by_x(result_list):
    '''
    按照起始坐标聚类
    :param result_list:
    :return:
    '''
    classfy_map = {}

    for index in range(0, len(result_list)):
        result_template = result_list[index]
        ptx = result_template.pt[0]

        if classfy_map.__contains__(ptx):
            classfy_map[ptx].append(result_template)
        elif classfy_map.__contains__(ptx - 1):  # 1像素容错
            classfy_map[ptx - 1].append(result_template)
        elif classfy_map.__contains__(ptx - 2):  # 2像素
            classfy_map[ptx - 2].append(result_template)
        else:
            group_list = [result_template]
            classfy_map[ptx] = group_list

    final_list = []  # 第一次优化，得到分组中匹配到的面积最大值，消除重叠

    def sort_area(tp):
        return tp.area

    for k in classfy_map:
        target_list = classfy_map[k]
        target_list.sort(key=sort_area, reverse=True)

        # for tp in target_list:
        #     print("分组坐标:", k, "标签:", tp.template.tempLabel, "实际坐标", tp.pt[0],"width ",tp.width, "面积:", tp.template.area())

        fe = target_list[0]  # 得到面积最大值
        final_list.append(fe)
        # print("最优标签:", fe.template.tempLabel, "x坐标：", fe.pt[0], "最大面积", fe.area, "width:", fe.width)
        # print("----------")

        # 第二次优化，消除X坐标重叠在一起的局部匹配结果
    left_delta = 2  # 前后两个验证码重叠值，这个重叠可能是局部重叠
    deleted_set = []
    for xfindex in range(1, len(final_list)):  # 从第1个开始
        this_template = final_list[xfindex]
        pre_tempate = final_list[xfindex - 1]
        if (pre_tempate.pt[0] + pre_tempate.width) - (this_template.pt[0]) >= left_delta:
            deleted_set.append(final_list[xfindex])
    # 得到分组终识别结果
    for needelete in deleted_set:
        final_list.remove(needelete)

    return final_list


def merge_result(match_results):
    '''
    :param match_results:
    :return:
    '''
    final_list = classfy_by_x(match_results)
    return final_list


def test():
    template_map = read_templates("./labeled")
    print("总模板类别:", len(template_map), "个")
    total_tempate_count = 0
    for tpk in template_map:
        total_tempate_count = total_tempate_count + len(template_map[tpk])
    print("总模板:", total_tempate_count, "个")

    threshold = 0.8  # 匹配阀值，超过这个值认为匹配正确
    root_dir = "./test"
    files = os.listdir(root_dir)

    total = 0
    right = 0

    for f in files:

        absfile = root_dir + os.sep + f
        if not absfile.endswith(".png"):
            continue
        total = total + 1
        # print("路径：", absfile)
        input_img = cv2.imread(absfile)

        cv2.threshold(input_img, 180, 255, cv2.THRESH_BINARY, input_img)
        input_img = filter_out_black(input_img)

        match_results = do_match(template_map, input_img, threshold)
        final_list = classfy_by_x(match_results)

        rec_result = ""
        print("-------------------------")
        for result in final_list:
            rec_result = rec_result + result.template.tempLabel;
            print("标签:", result.template.tempLabel, ",左上角坐标:", (result.pt[0], result.pt[1]), "w * h:",
                  (result.width, result.height))

        if f.lower().replace(".png", "") == rec_result.lower():
            right = right + 1
        else:
            print(absfile)
            print(rec_result)
            pass
    print(right, "/", total)
    print("-------------------------")


if __name__ == '__main__':

    is_test = 0

    if is_test != 1:
        template_map = read_templates("./labeled")
        print("总模板类别:", len(template_map), "个")
        total_tempate_count = 0
        for tpk in template_map:
            total_tempate_count = total_tempate_count + len(template_map[tpk])
        print("总模板:", total_tempate_count, "个")

        input_img = cv2.imread("./test/FqXV.png")

        cv2.threshold(input_img, 180, 255, cv2.THRESH_BINARY, input_img)
        input_img = filter_out_black(input_img)

        threshold = 0.8  # 匹配阀值，超过这个值认为匹配正确
        match_results = do_match(template_map, input_img, threshold)
        final_list = classfy_by_x(match_results)

        print("-------------------------")
        for result in final_list:
            print("标签:", result.template.tempLabel, ",左上角坐标:", (result.pt[0], result.pt[1]), "w * h:",
                  (result.width, result.height))
        print("-------------------------")
    else:
        test()

# all_img = cv2.imread('3h3m.png')
# all_img = cv2.cvtColor(all_img, cv2.COLOR_BGR2GRAY)
# template = cv2.imread('3.jpeg', 0)
# h, w = template.shape[:2]
#
# res = cv2.matchTemplate(all_img, template, cv2.TM_CCOEFF_NORMED)
# threshold = 0.7  # 取匹配程度大于%70的坐标
#
# loc = np.where(res >= threshold)
# for pt in zip(*loc[::-1]):  # *号表示可选参数
#     bottom_right = (pt[0] + w, pt[1] + h)
#     cv2.rectangle(all_img, pt, bottom_right, (0, 0, 255), 2)
#
# cv2.imshow('img_rgb', all_img)
# cv2.waitKey(0)
